'Handles connections to nodes.'
import logging
LOGGER = logging.getLogger('mapreduce.master.node')

import paramiko

from forest.connection import Connection

logging.getLogger('paramiko.transport').setLevel(logging.ERROR)

import tempfile

class SSHFile(object):
    def __init__(self, filename, connection, mode="r"):
        self.connection = connection
        self.filename = filename
        self.mode = mode
        self.tempfile = tempfile.NamedTemporaryFile()
        
        self.connection.get(self.filename, self.tempfile)
        self.fileobj = open(self.tempfile, mode)
        
    def read(self, *args, **kwargs):
        self.fileobj.read(*args, **kwargs)
        
    def __iter__(self, *args, **kwargs):
        self.fileobj.__iter__(*args, **kwargs)
    
    def readlines(self, *args, **kwargs):
        self.fileobj.readlines(*args, **kwargs)
        
    def write(self, *args, **kwargs):
        self.fileobj.write(*args, **kwargs)
        
    def close(self, *args, **kwargs):
        self.fileobj.close(*args, **kwargs)
        self.connection.put(self.tempfile, self.filename)
        
    def __getattr__(self, name):
        if name in ['connection', 'filename', 'fileobj', 'tempfile']:
            return self.__dict__.get(name)
            
        return self.fileobj.__dict__.get(name) 
    
    def seek(self, *args, **kwargs):
        self.fileobj.seek(*args, **kwargs)

class SSHConnection(Connection):
    'Essentially a ssh connection object to a worker. Like fabric.'
    
    channel = None
    
    def __init__(self, address):
        self.client = paramiko.SSHClient()
        self.client.load_system_host_keys()
        self.original_address = address
        self.uname, self.address = address.split('@')
        self.client.connect(self.address,  username=self.uname)
        self.transport = self.client.get_transport()
        #transport.set_keepalive(1)
     
    def debug(self, *args):
        args = list(args)
        args[0] = '[%s]: %s' % (self.original_address, args[0])
        LOGGER.debug(*args)
    
    def warn(self, *args):
        args = list(args)
        args[0] = '[%s]: %s' % (self.original_address, args[0])
        LOGGER.warn(*args)
        
    def get_channel(self):
        self.channel = self.transport.open_session()
        return self.channel
    
    def open(self, filename, mode="r"):
        ''' Returns a ssh file object that can be manipulated like a real file object.
        
            This is done by pulling the file down, and returning the open file handler.
            
            The filehandler is wrapped so whenever a write is called, the file
            is transferred back to the original location on the remote
        '''
        return SSHFile(filename, mode)

    def execute_wait(self, command, acceptable_codes=None, timeout=None):
        if acceptable_codes is None:
            acceptable_codes = []
        
        channel = self.get_channel()
        channel.settimeout(timeout)
        self.debug('%s' % command)
        channel.exec_command(command)
        code = self.wait()
        if code != 0 and code not in acceptable_codes:
            raise Exception('Command %s exited with %s' % (command, code))
        return code
        
    def execute(self, command, bufsize=-1, timeout=None):
        channel =self.get_channel()
        channel.settimeout(timeout)
        self.debug('%s' % command)
        channel.exec_command(command)
        stdin = channel.makefile('wb', bufsize)
        stdout = channel.makefile('rb', bufsize)
        stderr = channel.makefile_stderr('rb', bufsize)
        return stdin, stdout, stderr
        
    def send_data(self, data):
        self.channel.send(data)
        
    def wait(self):
        'wait for a return code'
        return self.channel.recv_exit_status()
        
    def readstderr(self):
        stderr_data = []
        while self.channel.recv_stderr_ready():
            stderr_data.append(self.channel.recv_stderr(1024))
        return "".join(stderr_data)
        
    def readstdout(self):
        stdout_data = []
        while self.channel.recv_ready():
            stdout_data.append(self.channel.recv(1024))
        return "".join(stdout_data)
    
    _sftp = None
    def sftpclient(self):
        if self._sftp is None:
            self._sftp = SFTPClient2.from_transport(self.transport)
        return self._sftp
    
        
    last_bytestotransfer = 0
    def last_bytes_transferred(self):
        return self.last_bytestotransfer
        
    def transfer_callback(self, bytestransferred, bytestotransfer):
        self.last_bytestotransfer = bytestotransfer
        
    def put(self, local_path, remote_path):
        self.debug('Putting %s to %s' % (local_path, remote_path))
        sftp = self.sftpclient()
        sftp.put(local_path, remote_path, callback=self.transfer_callback)
        
    def get(self, remote_path, local_path):
        self.debug('Getting %s to %s' % (remote_path, local_path))
        sftp = self.sftpclient()
        sftp.get(remote_path, local_path, callback=self.transfer_callback)
        
    def putfilelike(self, filelike, remote_path, filesize):
        self.debug('Putting buffer to %s' % (remote_path))
        sftp = self.sftpclient()
        return sftp.putfilelike(filelike, remote_path,  filesize, callback=self.transfer_callback)
        
    # common commands
    def read_file(self, path):
        try:
            self.execcommandandwait('cat '+path)
            return self.readstdout()
        except Exception, e:
            raise Exception("Error reading file %s: %s: %s" % (path, e, self.readstderr()))
        
    def exists(self, path, directory=False):
        command = 'test -e '+path
        code = self.execcommandandwait(command, acceptable_codes=[1])
        
        if code== 0:
            return True
            
        return False
    
import os

class SFTPClient2(paramiko.SFTPClient):
    def putfilelike(self, local_file, remotepath, file_size, callback=None, confirm=True):
        fl = local_file
        #file_size = os.fstat(fl.fileno()).st_size
        try:
            fr = self.file(remotepath, 'wb')
            fr.set_pipelined(True)
            size = 0
            try:
                while True:
                    data = fl.read(32768)
                    if len(data) == 0:
                        break
                    fr.write(data)
                    size += len(data)
                    if callback is not None:
                        callback(size, file_size)
            finally:
                fr.close()
        finally:
            fl.close()
        if confirm:
            s = self.stat(remotepath)
            if s.st_size != size:
                raise IOError('size mismatch in put!  %d != %d' % (s.st_size, size))
        else:
            s = SFTPAttributes()
        return s

    def getfilelike(self, remotepath, local_file, callback=None):
        fr = self.file(remotepath, 'rb')

